#include "pch.h"
#include "ThreadpoolTimer.h"
#include "AsyncLibImpl.h"

namespace AsyncLib
{

CThreadpoolTimer::CThreadpoolTimer()
{
	ZeroMemory(m_pTpCaches, sizeof(m_pTpCaches));
}

CThreadpoolTimer::~CThreadpoolTimer()
{
	assert(m_mapTasks.empty());
}

void CThreadpoolTimer::Init()
{
	m_mapTasks.reserve(ARRAYSIZE(m_pTpCaches));
}

void CThreadpoolTimer::Stop()
{
	std::unique_lock lock(m_mutex);

	for (auto& [pTpTimer, _] : m_mapTasks)
	{
		SetThreadpoolTimer(pTpTimer, NULL, 0, 0);
	}
}

void CThreadpoolTimer::Exit()
{
	// They have been waited and closed outside
	ZeroMemory(m_pTpCaches, sizeof(m_pTpCaches));
	m_mapTasks.clear();
}

void CThreadpoolTimer::AddTask(UINT nMilliseconds, std::function<void()>&& fnTask)
{
	auto pTpTimer = GetTpTimer();
	assert(pTpTimer);

	FILETIME ftDueTime;
	(INT64&)ftDueTime = -(nMilliseconds * 10000i64);

	std::unique_lock lock(m_mutex);
	m_mapTasks.emplace(pTpTimer, std::move(fnTask));
	SetThreadpoolTimer(pTpTimer, &ftDueTime, 0, 0);
}

void CThreadpoolTimer::OnCallback(PTP_TIMER pTpTimer)
{
	g_asyncLib.SetInThreadPool();
	std::unique_lock lock(m_mutex);

	if (g_asyncLib.IsRunning())
	{
		ReleaseTpTimer(pTpTimer);
		auto node = m_mapTasks.extract(pTpTimer);
		lock.unlock();
		assert(node);

		if (node)
			node.mapped()();
	}
}

PTP_TIMER CThreadpoolTimer::GetTpTimer()
{
	for (int i = 0; i != ARRAYSIZE(m_pTpCaches); i++)
	{
		auto pTpTimer = (PTP_TIMER)InterlockedExchangePointer((PVOID *)&m_pTpCaches[i], NULL);
		if (pTpTimer)
			return pTpTimer;
	}

	return CreateThreadpoolTimer([](PTP_CALLBACK_INSTANCE, PVOID pContext, PTP_TIMER pTpTimer)
	{
		((CThreadpoolTimer *)pContext)->OnCallback(pTpTimer);
	}, this, g_asyncLib.GetCallbackEnv());
}

void CThreadpoolTimer::ReleaseTpTimer(PTP_TIMER pTpTimer)
{
	for (int i = 0; i != ARRAYSIZE(m_pTpCaches); i++)
	{
		pTpTimer = (PTP_TIMER)InterlockedExchangePointer((PVOID *)&m_pTpCaches[i], pTpTimer);
		if (pTpTimer == NULL)
			return;
	}

	CloseThreadpoolTimer(pTpTimer);
}

} // namespace AsyncLib
